import c4d
import os

from c4d import gui, plugins, bitmaps
from c4d.modules import mograph as mo

PLUGIN_ID = 1037232

NONE_ID = -1
EXECUTE_ID =1000
DELAY_TIME_ID = 500
VERTEX_ANIMATION_PLUGIN_ID = 1037231

def Current_state_to_object(op, container=c4d.BaseContainer()):
    
    if not isinstance(op, c4d.BaseObject):
        raise TypeError('expected c4d.BaseObject, got %s.' %
                        op.__class__.__name__)

    doc = op.GetDocument()

    if not doc and (op.GetNext() or op.GetPred() or op.GetUp()):
        flags = c4d.COPYFLAGS_NO_ANIMATION | c4d.COPYFLAGS_NO_BITS | c4d.COPYFLAGS_NO_MATERIALPREVIEW
        op = op.GetClone(flags)

    if not doc:
        csto_doc = c4d.documents.BaseDocument()
        csto_doc.InsertObject(op)
    else:
        csto_doc = doc

    result = c4d.utils.SendModelingCommand(
        c4d.MCOMMAND_CURRENTSTATETOOBJECT, [op], c4d.MODELINGCOMMANDMODE_ALL,
        container, csto_doc)[0]

    if not doc:
        op.Remove()

    return result

def SingleKeyframeCover(doc, obj, position):

    TRACK_POS_X = "Position . X"
    TRACK_POS_Y = "Position . Y"
    TRACK_POS_Z = "Position . Z"
    
    TRACK_KEYWORD_ARR = [TRACK_POS_X, TRACK_POS_Y, TRACK_POS_Z]

    objTracks = obj.GetCTracks()
    if len(objTracks)==0 : return
    
    for track in objTracks:
    
        if not track.GetName() in TRACK_KEYWORD_ARR:continue;
    
        curve = track.GetCurve()
        keyDict = curve.FindKey(doc.GetTime())
        key = None
        if keyDict != None:
            key = keyDict["key"]

        if key == None:
            key = curve.AddKey(doc.GetTime())["key"]    #Current time
    
        if track.GetName() == TRACK_POS_X:
            key.SetValue(curve, position.x)
            
        if track.GetName() == TRACK_POS_Y:
            key.SetValue(curve, position.y)
            
        if track.GetName() == TRACK_POS_Z:
            key.SetValue(curve, position.z)

def FrameToTime(doc, frame):
    return c4d.BaseTime(float(frame)/doc.GetFps())

def TryCreateKeyframeTrack(obj):
    trPosX = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_POSITION, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_X, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trPosX) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trPosX))

    trPosY = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_POSITION, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_Y, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trPosY) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trPosY))

    trPosZ = c4d.DescID(c4d.DescLevel(c4d.ID_BASEOBJECT_POSITION, c4d.DTYPE_VECTOR, 0), c4d.DescLevel(c4d.VECTOR_Z, c4d.DTYPE_REAL, 0))
    if obj.FindCTrack(trPosZ) == None:
        obj.InsertTrackSorted(c4d.CTrack(obj, trPosZ))

def UpdateVertexHandle(doc, vertexList, targetIndex, targetValue):

    filterIndex = 0
    filterObject = None
    
    for item in vertexList:
        index,object = item
        
        if index == targetIndex:
            filterIndex = index
            filterObject = object

    TryCreateKeyframeTrack(object)
    SingleKeyframeCover(doc,filterObject,targetValue)

def IsContainVetex(collection, targetIndex):
    result = False
    
    for item in collection:
        i,baseObject = item
        
        if i == targetIndex:
            result = True
            break
            
    return result
    
def CreateDummyObject(doc, pointIndex, parentObject):

    newObject = c4d.BaseObject(c4d.Onull)
    newObject.SetName("Point"+str(pointIndex))
    
    r18_question = c4d.BaseObject(c4d.Onull)
    r18_question.SetName("_")
    
    newObject.InsertUnderLast(parentObject)
    r18_question.InsertUnderLast(newObject)
    
    return newObject

def WriteToKeyframe(doc, vetexHandleList, lastFrameVertexList, parentObject):
    currentFrameVertexList = list()
    activeObj = doc.GetActiveObject()
    collapseObject = Current_state_to_object(activeObj)
    allPoints = collapseObject.GetAllPoints()

    i = 0
    for item in allPoints:
    
        if len(lastFrameVertexList) > 0:
            index,value = lastFrameVertexList[i]
            
            if not IsContainVetex(vetexHandleList, i):
                newObj = CreateDummyObject(doc, i, parentObject)
                vetexHandleList.append((i, newObj))
                
            UpdateVertexHandle(doc, vetexHandleList, i, item)
    
        currentFrameVertexList.append((i,item))
        i=i+1

    collapseObject.Remove()
    
    return currentFrameVertexList
    
def ExecuteConvert(doc):
    lastFrameVetexList = list()
    vetexHandleList = list()

    activeObj = doc.GetActiveObject()

    if activeObj == None : return

    c4d.CallCommand(12501)#Go to start
    doc.ExecutePasses(None, True, True, True, c4d.BUILDFLAGS_0)
    collapseObject = Current_state_to_object(activeObj)

    i = 0
    for item in collapseObject.GetAllPoints():
        lastFrameVetexList.append((i,item))
        i = i+1

    groupObject = c4d.BaseObject(c4d.Onull)
    groupObject.SetName("Vetex_Handle_Group")
    
    doc.InsertObject(groupObject)
        
    fps = doc.GetFps()
    minTime = doc.GetMinTime()
    frameCount = minTime.GetFrame(fps)
    
    while True:

        if c4d.BaseTime(frameCount, fps) > doc.GetMaxTime():
            c4d.StatusClear()
            break

        doc.SetTime(c4d.BaseTime(frameCount, fps))
        
        executeResult = doc.ExecutePasses(None, True, True, True, c4d.BUILDFLAGS_0)

        if not executeResult:
            print "error exit"
            break
        
        lastFrameVetexList = WriteToKeyframe(doc, vetexHandleList, lastFrameVetexList, groupObject)

        frameCount = frameCount + 1
        
        #------------StatusSetBar
        currentTime = c4d.BaseTime(frameCount, fps).Get() - minTime.Get()
        maxTime = doc.GetMaxTime().Get() - minTime.Get()
        
        if currentTime <= 0: currentTime = 1
        if maxTime <= 0: maxTime = 1
        
        c4d.StatusSetBar((currentTime / maxTime) * 100)
        #------------------------
    
    activeObjUp = activeObj.GetUp()
    if not activeObjUp == None:
        collapseObject.InsertUnder(activeObjUp)

    doc.InsertObject(collapseObject)
    activeObj.InsertUnder(collapseObject)
    activeObj.SetEditorMode(c4d.MODE_OFF)
    activeObj.SetRenderMode(c4d.MODE_OFF)
    groupObject.InsertUnder(collapseObject)
    groupObject.InsertTag(c4d.BaseTag(VERTEX_ANIMATION_PLUGIN_ID))

class MyDialog(gui.GeDialog):

    def CreateLayout(self):
        self.SetTitle(title="Vertex Animation Baker")

        self.SetInt32(id=DELAY_TIME_ID, value=100)
        
        self.AddButton(id=EXECUTE_ID, flags=c4d.BFH_LEFT, initw=700, name="Execute")

        return True

    def Command(self, id, msg):
        
        doc = c4d.documents.GetActiveDocument()

        if id == EXECUTE_ID:
        
            ExecuteConvert(doc);
            
        c4d.EventAdd()
        
        return True

class Entry(plugins.CommandData):

    dialog = None

    def Execute(self, doc):
        if self.dialog is None:
           self.dialog = MyDialog()

        return self.dialog.Open(dlgtype=c4d.DLG_TYPE_ASYNC, pluginid=PLUGIN_ID)

    def RestoreLayout(self, sec_ref):
        if self.dialog is None:
           self.dialog = MyDialog()

        return self.dialog.Restore(pluginid=PLUGIN_ID, secret=sec_ref)

if __name__ == "__main__":
    bmp = bitmaps.BaseBitmap()
    dir, file = os.path.split(__file__)
    fn = os.path.join(dir, "res", "tag.png")
    bmp.InitWith(fn)
    plugins.RegisterCommandPlugin(id=PLUGIN_ID, str="VertexAnimationBaker",
                                info=0, icon=bmp, 
                                help="Vertex Animation Baker",
                                dat=Entry())